import numpy as np
import scipy as sp
import pandas as pd

import common
import ice


class Node:
    """
    Node of the graph representing all available routs.

    Attributes:
        index: Index of the node in the graph.
        name: Name of the node.
        lat: Latitude (deg).
        lon: Longitude (deg).
    """
    index: int
    name: str
    lat: float
    lon: float

    def __init__(self, index: int, name: str, lat: float, lon: float):
        """
        Args:
            index: Index of the node in the graph.
            name: Name of the node.
            lat: Latitude (deg).
            lon: Longitude (deg).
        """
        self.index = index
        self.name = name
        self.lat = lat
        self.lon = lon

    def __str__(self):
        return f"{self.name} ({self.lat:.2f}° N, {self.lon:.2f}° E)"

    def __eq__(self, other):
        return self.index == other.index


class Edge:
    """
    Edge of the graph representing all available routs.

    Attributes:
        node_start: Start node.
        node_end: End node.
        length: Length of the edge (nmi).
    """
    node_start: Node
    node_end: Node
    length: float

    def __init__(self, node_start: Node, node_end: Node, ice_cover: ice.IceCover):
        """
        Args:
            node_start: Start node.
            node_end: End node.
        """
        self.node_start = node_start
        self.node_end = node_end
        self.length = 0.0
        self.speed_60mw = []
        self.part = []

        # Build the related path.
        self.build_path(ice_cover)

    def __str__(self):
        return f"{self.node_start.name} --> {self.node_end.name} ({self.length:.2f} nmi)"

    def build_path(self, ice_cover: ice.IceCover):
        """
        Build the path related to this edge.
        """

        # Find the shortest path.
        path = ice_cover.find_path(self.node_start.lat, self.node_start.lon, self.node_end.lat, self.node_end.lon)

        # Save the results.
        self.part = path.part
        self.length = path.length
        self.speed_60mw = path.speed_60mw

    def compute_ice(self, part: float):
        """
        Compute the surrounded ice conditions given a passed part of the edge.

        Args:
            part: Part of the edge that has been passed.

        Returns:
            The speed of a 60 MW icebreaker (knots).
        """
        for i in range(1, len(self.part)):
            a = part - self.part[i - 1]
            b = self.part[i] - part
            c = self.part[i] - self.part[i - 1]
            if a >= 0 and b >= 0:
                return (self.speed_60mw[i - 1] * b + self.speed_60mw[i] * a) / c


class Graph:
    """
    Graph representing all available routs.

    Attributes:
        nodes: Nodes.
        edges: Edges.
    """
    nodes: list[Node]
    edges: list[Edge]

    def __init__(self, filepath: str, ice_cover: ice.IceCover):
        """
        Args:
            filepath: Path to a file with graph data.
        """

        # Create a list of nodes.
        self.nodes = []
        data_nodes = pd.read_excel(filepath, sheet_name="points")
        for i, row in data_nodes.iterrows():
            index = row["point_id"]
            lat = row["latitude"]
            lon = row["longitude"]
            name = row["point_name"]
            node = Node(index, name, lat, lon)
            self.nodes.append(node)

        # Create a list of edges.
        self.edges = []
        data_edges = pd.read_excel(filepath, sheet_name="edges")
        for i, row in data_edges.iterrows():
            i = row["start_point_id"]
            node_start = self.nodes[i]
            i = row["end_point_id"]
            node_end = self.nodes[i]
            edge = Edge(node_start, node_end, ice_cover)
            if np.isinf(edge.length):
                continue
            self.edges.append(edge)
            edge = Edge(node_end, node_start, ice_cover)
            if np.isinf(edge.length):
                continue
            self.edges.append(edge)

        # Compute the weight matrix.
        weights = np.zeros((len(self.nodes), len(self.nodes)))
        for edge in self.edges:
            weights[edge.node_start.index, edge.node_end.index] = edge.length

        # Find the shortest path for each pair of nodes.
        self.distances, self.predecessors =\
            sp.sparse.csgraph.shortest_path(weights, method="auto", directed=False, return_predecessors=True)

    def find_path(self, index_start: int, index_end: int) -> tuple[list[Edge], float]:
        """
        Find the shortest path between two nodes.

        Args:
            index_start: Index of the start node.
            index_end: Index of the end node.

        Returns:
            A tuple consisting of the shortest path (defined as a sequence of edges) and its length.
        """

        # Reconstruct the path.
        path_indices = [index_end]
        index = index_end
        while index != index_start and index != -9999:
            index = self.predecessors[index_start, index]
            path_indices.append(index)
        path_indices = list(reversed(path_indices))

        # Form the sequence of edges.
        path_edges = []
        for i in range(1, len(path_indices)):
            for edge in self.edges:
                if edge.node_start.index == path_indices[i - 1] and edge.node_end.index == path_indices[i]:
                    path_edges.append(edge)

        return path_edges, self.distances[index_start, index_end].flat[0]
